#include "stdafx.h"

#include <vector>
#include <unordered_map>

#include <WinSock2.h>

#include <gtest/gtest.h>
#include "../../../XProcess/JmpHooker.h"


#pragma comment(lib, "ws2_32")

using namespace XProcess;
using namespace std;

wchar_t addresstr[] = L"127.0.0.1:80";
char capturedText[1024];
vector<SOCKET> sockets;
std::unordered_map<SOCKET, SOCKADDR> connections;

SOCKET PASCAL FAR hook_socket(decltype(&socket) realSocket,
    _In_ int af,
    _In_ int type,
    _In_ int protocol)
{
    auto s = realSocket(af, type, protocol);
    sockets.push_back(s);
    return s;
}

int PASCAL FAR hook_connect(decltype(&connect) realConnect,
    _In_ SOCKET s,
    _In_reads_bytes_(namelen) const struct sockaddr FAR *name,
    _In_ int namelen)
{
    auto retv = realConnect(s, name, namelen);
    if (0 == retv)
    {
        connections.emplace(s, *name);
    }
    return retv;
}

//int PASCAL FAR hook_recv(decltype(&recv) real_recv,
//    _In_ SOCKET s,
//    _Out_writes_bytes_to_(len, return) __out_data_source(NETWORK) char FAR * buf,
//    _In_ int len,
//    _In_ int flags)
//{
//    int retv = recv(s, buf, len, flags);
//    //char *p = nullptr;
//    //if ((p = strstr(buf, "<?php")) != nullptr)
//    //{
//    //    buf[0] = ' ';
//    //}
//    return retv;
//}


auto addr = &socket;
TEST(Hook, DllExports)
{
    auto hooker1 = JmpHooker<1, decltype(&socket)>::hook((decltype(&socket))Module::load("ws2_32").procAddress("socket"), &hook_socket);
    auto hooker2 = JmpHooker<2, decltype(&connect)>::hook((decltype(&connect))Module::load("ws2_32").procAddress("connect"), &hook_connect);

    EXPECT_EQ(0, sockets.size());
    SOCKET s = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    EXPECT_EQ(1, sockets.size());
    EXPECT_EQ(s, sockets[0]);

    SOCKADDR addr;
    int addrlen = sizeof(addr);
    WSAStringToAddressW(addresstr, AF_INET, nullptr, &addr, &addrlen);

    EXPECT_EQ(0, connections.size());
    int retv = connect(s, &addr, addrlen);
    if (retv == 0)
    {
        EXPECT_EQ(1, connections.size());
        EXPECT_EQ(AF_INET, addr.sa_family);
    }
    else
    {
        EXPECT_EQ(0, connections.size());
        EXPECT_NE(AF_INET, addr.sa_family);
    }
    hooker1.unHook();
    hooker2.unHook();
}